from utils import file_util
from gensim.models import word2vec
import pickle
import time
import os

def createWord2vecModel(file_path, out_dir):
    filePath_list, label_list, content_list = file_util.gen_files_labels(file_path, out_dir)
    stopwords_path = r'./data/stop_words.txt'
    stopwords = file_util.load_stopWords(stopwords_path)
    tmp_w2vec_path,jieba_content_list = file_util.getVocabularyList(content_list, out_dir, stopwords)
    sentences = word2vec.Text8Corpus(tmp_w2vec_path)
    # 训练模型，部分参数如下
    model = word2vec.Word2Vec(sentences, size=100, hs=1, min_count=1, window=3)
    model.save(u'./data/w2vec.model')
    # 模型的预测
    print('-----------------分割线----------------------------')
    print('tmp_w2vec_path:', tmp_w2vec_path)
    print('jieba_content_list: ', jieba_content_list)

# 内容转索引矩阵
def content2IndexMat(file_path,stopwords_path,out_dir):
    stopwords = file_util.load_stopWords(stopwords_path)
    filePath_list, label_list, content_list = file_util.gen_files_labels(file_path, out_dir)
    jieba_content_list = file_util.cutNoSwdForjieba(content_list, stopwords)
    saver_data(jieba_content_list, label_list, 1000)

# interval 多少文章保存成一个pickle文件
def saver_data(content_codes, label_one_hots, interval):
    directory_name = 'content_list'
    if not os.path.isdir(directory_name):
        os.mkdir(directory_name)
    # 调用pickle库保存label_one_hots
    with open('data/label_one_hots.pickle', 'wb') as file:
        pickle.dump(label_one_hots, file)
    # 调用pickle库保存content_codes
    len_content = len(content_codes)
    startTime = time.time()
    for i in range(0, len_content, interval):
        startIndex = i
        endIndex = i + interval
        content_list = []
        print('%06d-%06d start' % (startIndex, endIndex))
        for content_code in content_codes[startIndex:endIndex]:
            content_list.append(content_code)
        save_fileName = directory_name + '/%06d-%06d.pickle' % (startIndex, endIndex)
        with open(save_fileName, 'wb') as file:
            pickle.dump(content_list, file)
        used_time = time.time() - startTime
        print('%06d-%06d used time: %.2f seconds' % (startIndex, endIndex, used_time))

def main():
    out_dir = r'./data'
    THUCNews_demopath = r'F:/技术资料/舆情食材/demo'
    stopwords_path = r'./resources/stop_words.txt'
    #createWord2vecModel(THUCNews_demopath, out_dir)
    content2IndexMat(THUCNews_demopath, stopwords_path, out_dir)


if __name__ == '__main__':
    main()